import json
import random
from collections import defaultdict
import sys


data_p = sys.argv[1]


def get_source(item):
    """Identify the source of an item from possible keys."""
    if "datasource" in item:
        return item["datasource"]
    elif "source" in item:
        return item["source"]
    elif "metadata" in item and isinstance(item["metadata"], dict):
        return item["metadata"].get("dataset")
    return "default"  # No source identified


# Load the JSON data
with open(data_p, "r") as f:
    data = json.load(f)

# Group items by their source
groups = defaultdict(list)
i = 0
iv = 0
for item in data:
    source = get_source(item)
    # print(len(item['conversations']))
    if "image" in item and isinstance(item["image"], list) and len(item["image"]) > 2:
        # if i < 6:
        # print(f'{item} image more than 1..')
        i += 1
    # if len(item["conversations"]) > 2:
    #     continue
        # is_multiple_image = False
        # for idx, conv in enumerate(item['conversations']):
        #     if idx > 1 and '<image>' in conv['value']:
        #         is_multiple_image = True

        # if is_multiple_image:
        #     print(f'{item} have multiple conversations!')
        #     continue
    if "image" not in item:
        # print(f'pure text: {item}, {len(item['conversations'][1]['value'])}')
        continue
    if "video" in item:
        iv += 1

    if "image" in item and isinstance(item["image"], list) and len(item["image"]) > 6:
    # if ("image" in item and isinstance(item["image"], list) and len(item["image"]) == 1) or  ("image" in item and isinstance(item["image"], str)):
        groups[source].append(item)


print(f"Got {i} multi-images.")
print(f"Got {iv} videos.")

# Seed for reproducibility
# random.seed(42)
filtered_data = []

# Process each group
for source, items in groups.items():
    # Shuffle to ensure random selection
    # if 'nuscenes' in source or 'RAVEN' in source:
    #     filtered_data.extend(items)
    # else:
    #     random.shuffle(items)

    #     # Calculate 50% of the group (at least 1 item if not empty)
    #     group_size = len(items)
    #     k = max(1, int(group_size * 1)) if group_size > 0 else 0

    #     print(f'{source} {len(items)} {(len(items)/len(data))*100:.2f}%')
    #     # Select items and add to result
    # filtered_data.extend(items[:k])
    # if "RAVEN" in source:
    #     continue
    if '3D-LLM_3_datasets' in source:
        continue
    items = [itm for itm in items if len(itm['conversations'][1]['value'].split(' ')) > 12]
    filtered_data.extend(items)

# for d in data:
#     if 'image' in d:
#         filtered_data.append(d)
#     else:
#         print(f'{d} not have image ')

# Save the filtered data
sub_d = data_p.replace(".json", f"_multi_{len(filtered_data)}.json")
with open(sub_d, "w") as f:
    json.dump(filtered_data, f, indent=2, ensure_ascii=False)
print(f"saved into: {sub_d}")
